{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1ede8c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import h5py\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.image as mpimg\n",
    "import os\n",
    "import os.path as osp\n",
    "import imageio.v2 as imageio\n",
    "import cv2\n",
    "\n",
    "from scipy.stats import describe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd3e4d32",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compare(depth, depth1):\n",
    "    mask = depth != depth1\n",
    "    if mask.sum() == 0:\n",
    "        return 0, 0\n",
    "    diff = np.abs(depth[mask].astype(np.int32) - depth1[mask])\n",
    "    return diff.max(), diff.sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "899a56cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import re\n",
    "import cv2\n",
    "\n",
    "\n",
    "def load_pfm_file(file_path):\n",
    "    with open(file_path, \"rb\") as file:\n",
    "        header = file.readline().decode(\"UTF-8\").strip()\n",
    "\n",
    "        if header == \"PF\":\n",
    "            is_color = True\n",
    "        elif header == \"Pf\":\n",
    "            is_color = False\n",
    "        else:\n",
    "            raise ValueError(\"The provided file is not a valid PFM file.\")\n",
    "\n",
    "        dimensions = re.match(r\"^(\\d+)\\s(\\d+)\\s$\", file.readline().decode(\"UTF-8\"))\n",
    "        if dimensions:\n",
    "            img_width, img_height = map(int, dimensions.groups())\n",
    "        else:\n",
    "            raise ValueError(\"Invalid PFM header format.\")\n",
    "\n",
    "        endian_scale = float(file.readline().decode(\"UTF-8\").strip())\n",
    "        if endian_scale < 0:\n",
    "            dtype = \"<f\"  # little-endian\n",
    "        else:\n",
    "            dtype = \">f\"  # big-endian\n",
    "\n",
    "        data_buffer = file.read()\n",
    "        img_data = np.frombuffer(data_buffer, dtype=dtype)\n",
    "\n",
    "        if is_color:\n",
    "            img_data = np.reshape(img_data, (img_height, img_width, 3))\n",
    "        else:\n",
    "            img_data = np.reshape(img_data, (img_height, img_width))\n",
    "\n",
    "        img_data = cv2.flip(img_data, 0)\n",
    "\n",
    "    return img_data\n",
    "def _load_pose(path, ret_44=False):\n",
    "    f = open(path)\n",
    "    RT = np.loadtxt(f, skiprows=1, max_rows=4, dtype=np.float32)\n",
    "    assert RT.shape == (4, 4)\n",
    "    RT = np.linalg.inv(RT)  # world2cam to cam2world\n",
    "\n",
    "    K = np.loadtxt(f, skiprows=2, max_rows=3, dtype=np.float32)\n",
    "    assert K.shape == (3, 3)\n",
    "\n",
    "    if ret_44:\n",
    "        return K, RT\n",
    "    return K, RT[:3, :3], RT[:3, 3]  # , depth_uint8_to_f32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1b97570",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_cam(campath):\n",
    "    \"\"\"\n",
    "    Reads camera parameters from a _cam.txt file.\n",
    "\n",
    "    Args:\n",
    "        file (file object): The opened camera file.\n",
    "\n",
    "    Returns:\n",
    "        tuple: A tuple containing the intrinsic and extrinsic matrices, and depth range.\n",
    "    \"\"\"\n",
    "    with open(campath) as file:\n",
    "        words = file.read().split()\n",
    "        \n",
    "        # Extrinsic parameters\n",
    "        extrinsic = np.array([float(w) for w in words[words.index('extrinsic') + 1: words.index('intrinsic')]]).reshape(4, 4)\n",
    "        \n",
    "        # Intrinsic parameters\n",
    "        intrinsic = np.array([float(w) for w in words[words.index('intrinsic') + 1: words.index('intrinsic') + 10]]).reshape(3, 3)\n",
    "        \n",
    "        # Depth range\n",
    "        depth_range = (float(words[words.index('intrinsic') + 10]), float(words[words.index('intrinsic') + 11]))\n",
    "        \n",
    "        return intrinsic, extrinsic, depth_range"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68b3c8d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "filename = '00000000_cam.txt'\n",
    "root = f\"/lc/data/3D/dtu/dtu_training/Cameras_train\"\n",
    "cam1 = osp.join(root, filename)\n",
    "# cam2 = osp.join(root, 'train', '00000000_cam.txt')\n",
    "k1, r1 = _load_pose(cam1, ret_44=True)\n",
    "print(k1, '\\n', r1)\n",
    "root = f\"/lc/data/3D/dtu/dtu_training/Cameras\"\n",
    "cam1 = osp.join(root, filename)\n",
    "# cam2 = osp.join(root, 'train', '00000000_cam.txt')\n",
    "k1, r1 = _load_pose(cam1, ret_44=True)\n",
    "print(k1, '\\n', r1)\n",
    "# intrinsic, extrinsic, depth_range = load_cam(cam1)\n",
    "# print(intrinsic, '\\n', extrinsic, '\\n', depth_range)\n",
    "# cam in train, are for depthmaps, scale 1/8 to 200X150, then center crop to 160X128"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f96de6ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "# root = f\"/lc/data/3D/dtu/dtu_training/Depths/scan9_train\"\n",
    "# path = f\"{root}/depth_map_0000.pfm\"\n",
    "# depth = load_pfm_file(path)\n",
    "# depth2 = cv2.imread(path.replace('map','visual').replace('pfm', 'png'), cv2.IMREAD_UNCHANGED)\n",
    "# print(depth.shape, depth2.shape)\n",
    "# print(describe(depth, axis=None))\n",
    "# print(describe(depth2, axis=None))\n",
    "# depth3 = depth2.astype(np.float32)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b00f3fcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "# Scannet++ depth dataset\n",
    "s = 0\n",
    "# List of file paths\n",
    "root = f\"/lc/data/3D/dtu/dtu_training/Depths/scan2_train\"\n",
    "file_paths = sorted([f for f in os.listdir(root) if f.endswith(\".pfm\")])\n",
    "\n",
    "cols = 2\n",
    "rows = 6\n",
    "interval = 1\n",
    "# Create a figure with 2 columns and 10 rows,dpi=100\n",
    "fig, axs = plt.subplots(rows, cols, figsize=(12, rows*3))\n",
    "\n",
    "# Iterate over the file paths and display the images\n",
    "for i, file_path in enumerate(file_paths[s:rows*cols*interval//2+s:interval]):\n",
    "    depth_file = osp.join(root, file_path)\n",
    "    depth = load_pfm_file(depth_file)\n",
    "    # depth2 = cv2.imread(depth_file.replace('map','visual').replace('pfm', 'png'), cv2.IMREAD_UNCHANGED)\n",
    "    id = int(file_path.split('_')[2].removesuffix('.pfm')) + 1\n",
    "    img_file = f'rect_{id:03d}_3_r5000.png'\n",
    "    img = mpimg.imread(osp.join(root.replace('Depths', 'Rectified'), img_file))\n",
    "    print(img.shape, depth.shape)\n",
    "    a = 3\n",
    "    b=6\n",
    "    axs[i * 2 // cols , i * 2 % cols ].imshow(img)\n",
    "    axs[i * 2 // cols , i * 2 % cols].axis(\"off\")  # Hide axis ticks\n",
    "    axs[i * 2 // cols , i * 2 % cols].set_title(img_file)\n",
    "    axs[i * 2 // cols , i * 2 % cols + 1].imshow(depth)\n",
    "    axs[i * 2 // cols , i * 2 % cols + 1].axis(\"off\")  # Hide axis ticks\n",
    "    axs[i * 2 // cols , i * 2 % cols + 1].set_title('depth')\n",
    "\n",
    "plt.show()\n",
    "# objects in various lighting conditions, each camera pose and depth map have 7 images\n",
    "# small objects\n",
    "# depth has low resolution\n",
    "# poses correct, no noise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5bfaf6f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "print(imgs[0].shape)\n",
    "print(depth.shape, depth.dtype) \n",
    "print(describe(depth, axis=None))\n",
    "# (512, 640, 3)\n",
    "# (128, 160) float32\n",
    "# DescribeResult(nobs=20480, minmax=(np.float32(0.0), np.float32(946.78094)), mean=np.float32(431.71103), variance=np.float64(104223.69060989306), skewness=np.float64(-0.42943912744522095), kurtosis=np.float32(-1.4890113))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2616181c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# extrinsic\n",
    "# 0.970263 0.00747983 0.241939 -191.02\n",
    "# -0.0147429 0.999493 0.0282234 3.28832\n",
    "# -0.241605 -0.030951 0.969881 22.5401\n",
    "# 0.0 0.0 0.0 1.0\n",
    "\n",
    "# intrinsic\n",
    "# 361.54125 0.0 82.900625\n",
    "# 0.0 360.3975 66.383875\n",
    "# 0.0 0.0 1.0\n",
    "\n",
    "# 425.0 2.5"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "vggt",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
